package com.basker.pisces.parallel;

import java.util.LinkedHashMap;
import java.util.Map;
import java.util.Map.Entry;
import java.util.concurrent.Callable;
import java.util.concurrent.ExecutionException;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Future;
import java.util.concurrent.SynchronousQueue;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;

import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.springframework.scheduling.concurrent.CustomizableThreadFactory;
import org.springframework.util.Assert;

import com.alibaba.ttl.TtlCallable;

/**
 * 并行执行多个任务的执行器
 *
 * <p>
 * <li>1.创建执行器；
 * <li>2.{@link #addTask(String, Callable)}添加任务；
 * <li>3.{@link #start()}启动任务；
 * <li>4.{@link #getResults()}或者{@link #getResult(String)}获取结果
 *
 * @author hangwen
 */
public class MultiTaskExecutor {

    private static final Logger LOGGER = LoggerFactory.getLogger(MultiTaskExecutor.class);

    private static final ExecutorService DEFAULT_EXECUTOR_SERVICE;

    static {
        ThreadFactory threadFactory = new CustomizableThreadFactory(
                MultiTaskExecutor.class.getSimpleName() + "-DefaultThread-");
        DEFAULT_EXECUTOR_SERVICE = new ThreadPoolExecutor(0, 128, 60L, TimeUnit.SECONDS,
                new SynchronousQueue<Runnable>(), threadFactory);
    }

    private String name = MultiTaskExecutor.class.getSimpleName();
    private ExecutorService executorService;

    private Map<String, Callable<?>> tasks = new LinkedHashMap<String, Callable<?>>();

    private Map<String, FutureWrapper> futures = new LinkedHashMap<String, FutureWrapper>();
    private AtomicBoolean started = new AtomicBoolean(false);

    public MultiTaskExecutor() {

    }

    /**
     * @param executorService 线程池对象，默认使用一个全局的CachedThreadPool
     */
    public MultiTaskExecutor(ExecutorService executorService) {
        this(MultiTaskExecutor.class.getSimpleName(), executorService);
    }

    /**
     * @param name 执行器的名称，默认就是MultiTaskExecutor
     */
    public MultiTaskExecutor(String name) {
        this(name, null);
    }

    /**
     * @param name            执行器的名称，默认就是MultiTaskExecutor
     * @param executorService 线程池对象，默认使用一个全局的CachedThreadPool
     */
    public MultiTaskExecutor(String name, ExecutorService executorService) {
        this.name = name;
        this.executorService = executorService;
    }

    /**
     * 添加任务
     *
     * @param taskName 任务的唯一标识，后续可以通过这个名称获取任务对应的结果
     * @param task
     */
    public void addTask(String taskName, Callable<?> task) {
        this.assertTaskName(taskName);
        Assert.notNull(task, "parameter 'task' is required");

        this.tasks.put(taskName, TtlCallable.get(task));
    }

    /**
     * 获取导致任务出错的异常对象。如果任务尚在执行中，会阻塞当前线程。
     *
     * @param taskName
     * @return
     */
    public Throwable getCause(String taskName) {
        this.assertTaskName(taskName);
        this.assertStarted();

        FutureWrapper futureWrapper = findFutureWrapper(taskName);
        return futureWrapper.isOk() ? null : futureWrapper.cause;
    }

    /**
     * 获取指定任务的结果，如果任务执行出错，返回null。如果任务尚在执行中，会阻塞当前线程。
     *
     * @param taskName
     * @return
     */
    public <T> T getResult(String taskName) {
        this.assertTaskName(taskName);
        this.assertStarted();

        FutureWrapper futureWrapper = findFutureWrapper(taskName);
        return futureWrapper.safeGet();
    }

    /**
     * 获取所有任务的执行结果，如果任务执行出错，对应任务的结果为null。如果任务尚在执行中，会阻塞当前线程。
     *
     * @return
     */
    public Map<String, ?> getResults() {
        this.assertStarted();

        Map<String, Object> results = new LinkedHashMap<>(this.futures.size(), 1);

        for (Entry<String, FutureWrapper> entry : this.futures.entrySet()) {
            String taskName = entry.getKey();
            FutureWrapper futureWrapper = entry.getValue();

            results.put(taskName, futureWrapper.safeGet());
        }

        return results;
    }

    /**
     * 判断所有任务是否都执行完成
     *
     * @return
     */
    public boolean isDone() {
        this.assertStarted();

        return this.futures.values().stream().allMatch(f -> f.future.isDone());
    }

    /**
     * 指定任务是否正常执行完毕，即没有出现异常。如果任务尚在执行中，会阻塞当前线程。
     *
     * @param taskName
     * @return
     */
    public boolean isOk(String taskName) {
        this.assertTaskName(taskName);
        this.assertStarted();

        FutureWrapper futureWrapper = findFutureWrapper(taskName);
        return futureWrapper.isOk();
    }

    /**
     * 开始执行任务，只能执行一次。
     */
    public void start() {
        if (this.tasks.isEmpty()) {
            return;
        }

        if (!started.compareAndSet(false, true)) {
            throw new RuntimeException(name + " has started!");
        }

        ExecutorService workers = this.getExecutorService();
        for (Entry<String, Callable<?>> entry : tasks.entrySet()) {
            String taskName = entry.getKey();
            Callable<?> task = entry.getValue();

            Future<?> future = workers.submit(task);

            this.futures.put(taskName, new FutureWrapper(taskName, future));
        }
    }

    /**
     * 阻塞当前线程，直到所有任务都执行完毕
     *
     * @return 如果任务都正常完成（不出错），返回true，否则false
     */
    public boolean waitForTasks() {
        this.assertStarted();
        return this.futures.values().stream().allMatch(f -> f.isOk());
    }

    protected String getName() {
        return name;
    }

    protected void setExecutorService(ExecutorService executorService) {
        this.executorService = executorService;
    }

    protected void setName(String name) {
        this.name = name;
    }

    private void assertStarted() {
        Assert.isTrue(this.started.get(), name + " does not invoke \"start\" method!");
    }

    private void assertTaskName(String taskName) {
        Assert.hasText(taskName, "parameter 'taskName' is required");
    }

    private FutureWrapper findFutureWrapper(String taskName) {
        FutureWrapper futureWrapper = this.futures.get(taskName);

        if (futureWrapper == null) {
            throw new IllegalArgumentException(taskName + " is not exist in " + name);
        }

        return futureWrapper;
    }

    private ExecutorService getExecutorService() {
        return this.executorService != null ? this.executorService : DEFAULT_EXECUTOR_SERVICE;
    }

    private static class FutureWrapper {
        private String taskName;
        private Future<?> future;
        private Object ret;
        private Throwable cause;

        public FutureWrapper(String taskName, Future<?> future) {
            this.taskName = taskName;
            this.future = future;
        }

        public boolean isOk() {
            this.safeGet();
            return this.cause == null;
        }

        @SuppressWarnings("unchecked")
        public <T> T safeGet() {
            if (this.ret != null) {
                return (T) this.ret;
            }

            if (this.cause != null) {
                return null;
            }

            try {
                this.ret = this.future.get();
                return (T) this.ret;
            } catch (InterruptedException e) {
                this.cause = e;
                Thread.currentThread().interrupt();

                return null;
            } catch (ExecutionException e) {
                this.cause = e.getCause();

                if (LOGGER.isErrorEnabled()) {
                    LOGGER.error(this.taskName + " occured an exception:" + cause.getMessage(), cause);
                }

                return null;
            }
        }
    }

}
